"""
@Filename       : iedp.py
@Create Time    : 2021/1/12 9:00
@Author         : Rylynn
@Description    : A Novel Embedding Method for Information Diffusion Prediction in Social Network Big Data
                IEEE TRANSACTIONS ON INDUSTRIAL INFORMATICS, VOL. 13, NO. 4, AUGUST 2017

"""

import torch as th
import torch.nn as nn


class IEDP(nn.Module):
    def __int__(self, config):
        super(IEDP, self).__int__()
        self.user_embed_s = nn.Embedding(config['user_num'], config['embed_size'])
        self.user_embed_i = nn.Embedding(config['user_num'], config['embed_size'])
        self.user_embed_l = nn.Embedding(config['user_num'], config['embed_size'])

        self.content_embed = nn.Embedding(config['contetn_num'], config['embed_size'])
        self.content_ws = nn.Linear(config['embed_size'], config['embed_size'], bias=False)
        self.content_wi = nn.Linear(config['embed_size'], config['embed_size'], bias=False)
        self.content_wl = nn.Linear(config['embed_size'], config['embed_size'], bias=False)

        self.margin = config['margin']

        nn.init.xavier_normal_(self.user_embed_s.weight)
        nn.init.xavier_normal_(self.user_embed_s.weight)
        nn.init.xavier_normal_(self.user_embed_l.weight)
        nn.init.xavier_normal_(self.content_embed.weight)
        nn.init.xavier_normal_(self.content_ws.weight)
        nn.init.xavier_normal_(self.content_wi.weight)
        nn.init.xavier_normal_(self.content_wl.weight)

    def forward(self, head, content, tail, is_neg):
        head_embed = self.user_embed_s(head)
        if is_neg:
            tail_embed = self.user_embed_l(tail)
        else:
            tail_embed = self.user_embed_i(tail)
        content_embed = self.content_embed(content)
        energy = self.content_ws(head_embed) + content_embed - self.content_wi(tail_embed)

        return th.norm(energy, self.norm)

    def loss(self, energy_pos, energy_neg):
        return th.sum(th.clamp_min(self.margin + energy_pos - energy_neg, 0))


def train_iedp():
    pass
